import io
import json
import paramiko
from threading import Thread
from channels.generic.websocket import WebsocketConsumer

from cmdb.backends.sshargs import args


class SSHBridge(object):
    def __init__(self, websocket):
        self.websocket = websocket

    def connect(self, host, port, username, authtype, password=None, pkey=None, term='xterm', cols=80, rows=24):
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        try:
            if authtype == 2:
                pkey = paramiko.RSAKey.from_private_key(io.StringIO(pkey))
                ssh.connect(username=username, hostname=host, port=port, pkey=pkey, timeout=8)
            else:
                ssh.connect(hostname=host, port=port, username=username, password=password, timeout=8)
        except Exception as e:
            message = json.dumps({'flag': 'error', 'message': str(e)})
            self.websocket.send(message)

            return False

        # 打开一个ssh通道并建立连接
        transport = ssh.get_transport()
        self.ssh_channel = transport.open_session()
        self.ssh_channel.get_pty(term=term, width=cols, height=rows)
        self.ssh_channel.invoke_shell()

        # 连接建立一次，之后交互数据不会再进入该方法
        for i in range(2):
            recv = self.ssh_channel.recv(1024).decode('utf-8','ignore')
            message = json.dumps({'flag': 'success', 'message': recv})
            self.websocket.send(message)

    def close(self):
        try:
            self.websocket.close()
            self.ssh_channel.close()
        except BaseException as e:
            pass

    def _ws_to_ssh(self, data):
        try:
            self.ssh_channel.send(data)
        except OSError as e:
            self.close()

    def _ssh_to_ws(self):
        try:
            while not self.ssh_channel.exit_status_ready():
                data = self.ssh_channel.recv(1024).decode('utf-8','ignore')
                if len(data) != 0:
                    message = {'flag': 'success', 'message': data}
                    self.websocket.send(json.dumps(message))
                else:
                    break
        except Exception as e:
            message = {'flag': 'error', 'message': str(e)}
            self.websocket.send(json.dumps(message))

            self.close()

    def shell(self, data):
        Thread(target=self._ws_to_ssh, args=(data,)).start()
        Thread(target=self._ssh_to_ws).start()

    def resize_pty(self, cols, rows):
        self.ssh_channel.resize_pty(width=cols, height=rows)


class SSHConsumer(WebsocketConsumer):
    def connect(self):
        self.pk = self.scope['url_route']['kwargs'].get('id')
        self.query = self.scope.get('query_string')
        self.user = self.scope['user']

        self.accept()

        # WebSocket连接成功后，连接ssh
        ssh_connect_args = args(self.pk, self.user, self.query)

        self.ssh = SSHBridge(websocket=self)
        self.ssh.connect(**ssh_connect_args)

    def disconnect(self, close_code):
        self.ssh.close()

    def receive(self, text_data=None):
        text_data = json.loads(text_data)

        if text_data.get('flag') == 'resize':
            self.ssh.resize_pty(cols=text_data['cols'], rows=text_data['rows'])
        else:
            self.ssh.shell(data=text_data.get('data', ''))
